# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment1/question4/fourier_transform.py
import torch

_i = torch.complex(
    torch.tensor(0, dtype=torch.float64),
    torch.tensor(1, dtype=torch.float64))


def discrete_fourier_transform(values_complex_vector: torch.Tensor):
    if len(values_complex_vector.shape) != 1:
        raise ValueError("values should be a vector.")
    n = values_complex_vector.shape[0]
    alpha = torch.exp(-_i * 2 * torch.pi / n)
    a = torch.arange(0, n)
    a = torch.reshape(a, [1, n]) * torch.reshape(a, [n, 1])
    a = alpha ** a
    return torch.matmul(a, values_complex_vector)
